import numpy as np
import rasterio
from rasterio.windows import Window
from scipy.stats import pearsonr
from tqdm import tqdm  # 导入 tqdm 库


def calculate_correlation(grid1, grid2, nodata1, nodata2):
    """
    计算两个网格的 Pearson 相关性，忽略 nodata 和 0 值。
    如果有效数据点不足10个，返回 NaN。
    """
    # 创建掩码，忽略 nodata 和 0 值
    mask1 = (grid1 != nodata1) & (grid1 != 0)
    mask2 = (grid2 != nodata2) & (grid2 != 0)
    valid_mask = mask1 & mask2

    # 检查有效数据点数量
    valid_count = np.sum(valid_mask)
    if valid_count < 10:
        return np.nan  # Pearson要求至少2个点，这里放了10个

    # 提取有效数据
    valid_data1 = grid1[valid_mask]
    valid_data2 = grid2[valid_mask]

    # 计算 Pearson 相关性
    correlation, _ = pearsonr(valid_data1, valid_data2)
    return correlation


def grid_correlation(tiff1_path, tiff2_path, output_path, grid_size):
    """
    计算两个 TIFF 图层的网格相关性，并导出结果。
    """
    with rasterio.open(tiff1_path) as src1, rasterio.open(tiff2_path) as src2:
        # 确保两个图层的行列数一致
        assert src1.width == src2.width and src1.height == src2.height, "两个图层的行列数不一致"

        # 获取图层的元数据
        meta = src1.meta.copy()
        meta.update(dtype=rasterio.float32, count=1, nodata=np.nan)

        # 获取图层的宽度和高度
        width = src1.width
        height = src1.height

        # 计算网格的行列数
        rows = int(np.ceil(height / grid_size))
        cols = int(np.ceil(width / grid_size))

        # 创建输出数组
        correlation_grid = np.zeros((rows, cols), dtype=np.float32)

        # 使用 tqdm 显示进度条
        for i in tqdm(range(rows), desc="Processing rows", unit="row"):
            for j in range(cols):
                # 计算当前网格的窗口
                row_start = i * grid_size
                col_start = j * grid_size
                row_end = min(row_start + grid_size, height)
                col_end = min(col_start + grid_size, width)

                window = Window(col_start, row_start, col_end - col_start, row_end - row_start)

                # 读取当前网格的数据
                grid1 = src1.read(1, window=window)
                grid2 = src2.read(1, window=window)

                # 计算相关性
                correlation = calculate_correlation(grid1, grid2, src1.nodata, src2.nodata)
                correlation_grid[i, j] = correlation

        # 导出结果
        with rasterio.open(output_path, 'w', **meta) as dst:
            dst.write(correlation_grid, 1)


# 示例调用
tiff1_path = r'D:\GPPvsWater\work\NDVImean_con_1km.tif'
tiff2_path = r'D:\GPPvsWater\work\ai_v3_yr_con_1km.tif'
output_path = r'D:\GPPvsWater\work\output_correlation2.tif'
grid_size = 25  # 网格大小

grid_correlation(tiff1_path, tiff2_path, output_path, grid_size)
